# src/image_processing/image_handler.py

import asyncio
from PIL import Image, UnidentifiedImageError  # 导入PIL库中的Image模块和UnidentifiedImageError异常
import numpy as np  # 导入numpy库，用于处理图像数组
from src.crypto.rc4 import RC4  # 导入RC4加密算法
from src.crypto.lfsr import LFSR  # 导入LFSR加密算法
import os  # 导入os库，用于处理文件路径
from werkzeug.utils import secure_filename  # 用于安全地处理文件名
import uuid  # 用于生成唯一文件名

class ImageData:
    def __init__(self, image_path=None, image_array=None, mode=None):
        """
        初始化图像类。
        :param image_path: 图像路径（可选）
        :param image_array: 图像数组（可选）
        :param mode: 图像模式（可选）
        """
        self.image_path = image_path  # 图像文件路径
        self.image_array = image_array  # 图像数组（numpy数组）
        self.mode = mode  # 图像模式（如 'RGB', 'L', 'RGBA' 等）

        # 如果提供了图像路径但没有提供图像数组，则加载图像
        if image_path is not None and image_array is None:
            self.load_image(image_path)

    def load_image(self, image_path):
        """
        从文件路径加载图像，并验证是否为有效图像。
        :param image_path: 图像路径
        :raises ValueError: 如果文件不是有效的图像文件
        """
        try:
            # 尝试打开图像文件
            img = Image.open(image_path)
            img.verify()  # 验证文件是否为有效图像
            img = Image.open(image_path)  # 重新打开文件以进行后续操作
            self.mode = img.mode  # 获取图像模式

            # 如果图像是 P 模式（调色板模式），则转换为 L、RGB 或 RGBA 模式
            if self.mode == 'P':
                # 生成转换后的文件名，确保支持中文
                base_name = os.path.splitext(os.path.basename(image_path))[0]
                converted_path = os.path.join(os.path.dirname(image_path), f"{base_name}_converted.png")
                self.mode = self.convert_palette_image(image_path, converted_path)  # 转换图像模式
                img = Image.open(converted_path)  # 打开转换后的图像
                self.image_array = np.array(img)  # 将图像转换为numpy数组
            else:
                self.image_array = np.array(img)  # 将图像转换为numpy数组
        except UnidentifiedImageError:
            raise ValueError("无法识别图像文件格式，请上传有效的图像文件。")
        except Exception as e:
            raise ValueError(f"加载图像时出错: {str(e)}")

    def convert_palette_image(self, image_path, output_path):
        """
        将 P 模式的图像转换为适合的类型（L、RGB 或 RGBA），并保存转换后的图像。
        :param image_path: 输入图像路径
        :param output_path: 输出图像路径
        :return: 转换后的图像模式（'L'、'RGB' 或 'RGBA'）
        """
        img = Image.open(image_path)
        if img.mode != 'P':
            raise ValueError("输入图像不是 P 模式（调色板模式）！")

        palette = img.getpalette()  # 获取图像的调色板
        if palette is None:
            raise ValueError("无法获取图像的调色板！")

        has_alpha = len(palette) == 256 * 4  # 调色板包含 RGBA 值
        is_grayscale = True  # 假设图像是灰度图

        # 检查调色板中的颜色是否为灰度
        if has_alpha:
            for i in range(0, len(palette), 4):
                r, g, b = palette[i], palette[i + 1], palette[i + 2]
                if r != g or g != b:
                    is_grayscale = False
                    break
        else:
            for i in range(0, len(palette), 3):
                r, g, b = palette[i], palette[i + 1], palette[i + 2]
                if r != g or g != b:
                    is_grayscale = False
                    break

        # 根据调色板内容决定转换模式
        if is_grayscale:
            converted_mode = 'L'  # 灰度模式
            img = img.convert('L')
        elif has_alpha:
            converted_mode = 'RGBA'  # 带透明度的RGB模式
            img = img.convert('RGBA')
        else:
            converted_mode = 'RGB'  # 普通RGB模式
            img = img.convert('RGB')

        img.save(output_path)  # 保存转换后的图像
        return converted_mode  # 返回转换后的图像模式

    def save_image(self, output_path):
        """
        保存图像到指定路径。
        :param output_path: 输出路径
        """
        if self.image_array is None:
            raise ValueError("图像数组为空，无法保存图像！")

        # 获取原始图像的格式
        if self.image_path:
            file_extension = os.path.splitext(self.image_path)[1].lower()
            if file_extension in ['.jpg', '.jpeg', '.png']:
                format = file_extension[1:].upper()  # 去掉点号并转换为大写
            else:
                format = 'PNG'  # 默认保存为 PNG 格式
        else:
            format = 'PNG'  # 如果没有原始路径，默认保存为 PNG 格式

        # 确保输出路径的目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        # 保存图像时指定格式，并确保支持中文文件名
        Image.fromarray(self.image_array).save(output_path, format=format)

    def get_image_array(self):
        """
        获取图像数组。
        :return: 图像数组
        """
        return self.image_array

    def get_mode(self):
        """
        获取图像模式。
        :return: 图像模式
        """
        return self.mode

    def get_image_path(self):
        """
        获取图像路径。
        :return: 图像路径
        """
        return self.image_path

    def encrypt(self, key, algorithm_name, output_path):
        """
        使用指定的算法加密图像。
        :param key: 加密密钥
        :param algorithm_name: 加密算法名称（如 "RC4" 或 "LFSR"）
        :param output_path: 加密后图像的保存路径
        :return: 加密后的 ImageData 对象
        """
        if algorithm_name.lower() == "rc4":
            # 将二进制密钥转换为字节
            rc4_key = bytes(int(key[i:i + 8], 2) for i in range(0, len(key), 8))
            rc4 = RC4(rc4_key)  # 初始化RC4加密器
            key_stream = rc4.encrypt(b"\x00" * self.image_array.size)  # 生成密钥流
        elif algorithm_name.lower() == "lfsr":
            seed = int(key, 2)  # 将二进制密钥转换为整数
            lfsr = LFSR(seed=seed, taps=[21, 2], num_bits=21)  # 初始化LFSR加密器
            key_stream = lfsr.generate_key_stream(self.image_array.size)  # 生成密钥流
        else:
            raise ValueError("不支持的加密算法！")

        # 将密钥流转换为numpy数组，并调整形状以匹配图像数组
        key_stream = np.frombuffer(key_stream, dtype=np.uint8)
        key_stream = np.resize(key_stream, self.image_array.shape)

        # 使用异或操作加密图像
        encrypted_array = np.bitwise_xor(self.image_array, key_stream)
        encrypted_image = Image.fromarray(encrypted_array)  # 将加密后的数组转换为图像

        # 获取原始图像的格式
        if self.image_path:
            file_extension = self.image_path.split('.')[-1].lower()
            if file_extension in ['jpg', 'jpeg', 'png']:
                format = file_extension.upper()
            else:
                format = 'PNG'  # 默认保存为 PNG 格式
        else:
            format = 'PNG'  # 如果没有原始路径，默认保存为 PNG 格式

        # 保存加密后的图像
        encrypted_image.save(output_path, format=format)

        return ImageData(image_path=output_path)  # 返回加密后的ImageData对象

def read_image(image_path):
    """
    读取图像并返回图像数组和模式。
    :param image_path: 图像路径
    :return: 图像数组和模式
    """
    image_data = ImageData(image_path=image_path)  # 初始化ImageData对象
    return image_data.get_image_array(), image_data.get_mode()  # 返回图像数组和模式

async def save_file_async(file, filepath):
    """
    异步保存文件
    :param file: 文件对象
    :param filepath: 文件保存路径
    """
    loop = asyncio.get_event_loop()  # 获取事件循环
    await loop.run_in_executor(None, file.save, filepath)  # 异步保存文件

def generate_unique_filename(original_filename, prefix='', suffix=''):
    """
    生成唯一的文件名，支持中文名称，并确保后缀合法。
    :param original_filename: 原始文件名
    :param prefix: 文件名前缀
    :param suffix: 文件名后缀
    :return: 生成的文件名
    """
    # 获取文件名和后缀
    base_name, ext = os.path.splitext(original_filename)
    # 确保后缀合法
    ext = ext.lower() if ext.lower() in ['.png', '.jpg', '.jpeg'] else '.png'
    # 生成新的文件名，使用 uuid 确保唯一性
    unique_name = f"{prefix}{base_name}_{uuid.uuid4().hex}{suffix}{ext}"
    return secure_filename(unique_name)